import os
import torch
from torchvision.utils import save_image
from tqdm import tqdm
from picanet_resnet18 import PiCANetResnet18

def test(model, dataloader, device,result_dir="results"):
    model.eval()
    os.makedirs(result_dir, exist_ok=True)
    with torch.no_grad():
        for idx, (images, _) in enumerate(tqdm(dataloader, desc="Testing")):
            images = images.to(device)
            _, _, _, x4 = model(images)
            pred = torch.nn.Conv2d(512, 1, kernel_size=1).to(device)(x4)
            pred = torch.nn.functional.interpolate(pred, size=(224, 224), mode='bilinear', align_corners=False)
            pred = torch.sigmoid(pred)

            for b in range(images.size(0)):
                save_path = os.path.join(result_dir, f"{idx * dataloader.batch_size + b:04d}.png")
                save_image(pred[b], save_path)
